import psqloperation as psql
import jsbeautify
import rmcomment
import chunkJS
import malurl
import random
import numpy as np

import commands
import sys
import os
import fnmatch

# For path_checksum
import hashlib
from os.path import normpath, walk, isdir, isfile, dirname, basename, \
	exists as path_exists, join as path_join


def checkandaskdir(outdir):
	'''Return True if output dir exists or user authorize to create, \
		o.w. False.'''
	if os.path.exists(outdir):
		return True
	else:
		decision = raw_input("{0} doesn't exist. Do you want to "
				"create [y/n]".format(outdir))
		if decision[0].lower() == 'y':
			os.makedirs(outdir)
			return True
	return False


def checkdir(inputdir, outputdir):
	'''Return False if input dir doesn't exist, o.w. True.'''
	if not os.path.exists(inputdir):
		print "Error: Input directory doesn't exist."
		return False
	if not os.path.exists(outputdir):
		os.makedirs(outputdir)	
	return True


def checkfile(inputfile, outputfile):
	'''Return False if the input is not a file or output already exists, \
		o.w. True.'''
	if not os.path.isfile(inputfile):
		return False
	if os.path.exists(outputfile):
		return False
	return True


def lsresult(regex):
	res = commands.getoutput('ls ' + regex)
	return res.split('\n')


def lsgrep(rels, regrep):
	res = commands.getoutput('ls ' + rels + ' | grep ' + regrep)
	return res.split('\n')

def ls2file(regex, outfile):
	commands.getoutput('ls -d ' + regex + ' > ' + outfile)


def genTestName(url):
	url_hash = hashlib.sha256(url).hexdigest()
	label = 'B1'	# means unknown 
	return genTempnameTuple(url, url_hash, label, True)


def genTempnameLine(line):
	line = line.split('\',\'')
	url = line[1]
	url_hash = line[0][2:]
	label = line[3]
	return genTempnameTuple(url, url_hash, label)


def genTempnameTuple(url, url_hash, label, test=False):
	outfilename = label + '.' + url_hash
	jsWithComment = '../../data/TempData/JSFileWComments/'
	jsNoComment = '../../data/TempData/JSFiles/'
	jsChunk = '../../data/TempData/JSFiles/'
	if test:
		jsBeautify = '../../data/JSTest/'
	else:
		jsBeautify = '../../data/JSAdd2Tree/'
	return url, jsWithComment+outfilename+'.js', jsNoComment+outfilename+'.js', jsChunk+outfilename, jsBeautify+outfilename 


def findHash(urltofind, cmpfile = '../../data/TempData/URLs/terryurl'):
	urltofind = 'http://ruian.gtisc.gatech.edu/Terry/' + urltofind
	cmpf = open(cmpfile, 'r')
	for compare in cmpf:
		line = compare.split('\',\'')
		url = line[1]
		url_hash = line[0][2:]
		if url == urltofind:
			return url_hash



def checkEmpty(result):
	'''
	Check whether one file leads to empty result.
	Compare the empty hashed filename with the record in URLs/terryurl.
	'''
	infile = '../../data/JSZipped/emptyfile'
	inf = open(infile, 'r')
	hashtofind = findHash(result)
	for line in inf:
		if line.split('.')[1][:-1] == hashtofind:
			return True
	return False




def D8Terry():
	'''
	Use modified d8 to run javascript given by Terry.
	'''
	indir = '../../data/JSTerry/javascript_20131018/'
	outdir = '../../data/TempData/JSFileWComments/'
	jsChunk = '../../data/TempData/JSFiles/'
	jsBeautify = '../../data/JSAdd2Tree/'

	results = lsresult(indir)
	for result in results:
		if not checkEmpty(result):
			continue
		# First
		# Use d8 to generate output
		# Remove the ones that have syntax error.
		'''
		commands.getoutput('d8 ' + indir + result + ' > ' + outdir + 'MTerry.' + findHash(result) + '.js')
		'''
		outfilename = 'MTerry.' + findHash(result)
		jscommentname = outdir + outfilename + '.js'
		if not os.path.exists(jscommentname):
			continue

		# Third
		# for scripts that doesn't work with modified chrome and v8
		# just beautify them
		print jscommentname
		beautifydir = jsBeautify + outfilename
		jsbeautify.jsdeobfuscatefile(indir + result, beautifydir + '/0')
		print beautifydir
	
		# Second
		# Use d8 to generate output
		# work on the modified v8 generated scripts
		'''
		chunkoutdir = jsChunk + outfilename
		beautifydir = jsBeautify + outfilename
		chunkJS.chunkJSFile(jscommentname, chunkoutdir)
		jsbeautify.jsdeobfuscatedir(chunkoutdir, beautifydir)
		'''
	

def NonsenseJSFromTerry(outfile = '../../data/JSTerry/javascript_20131018/NonsenseJSFromTerry'):
	indir = '../../data/JSTerry/javascript_20131018/'
	results = lsresult(indir)
	outf = open(outfile, 'w')
	for result in results:
		if checkEmpty(result):
			outf.write(result + '\n')
	outf.close()



class DictFile:
	"""
	Class to handle write dict to file and read dict from file operations.
	"""
	def dict2file(self, tostore, outfile):
		outf = open(outfile, 'w')
		for key, value in tostore.items():
			outf.write(str(key) + ',' + str(value) + '\n')
		outf.close()

	def file2dict(self, infile, valuetype):
		# If there is no such file, just return a new dict
		if not os.path.exists(infile):
			return dict()

		# Read the file to dict based on valuetype
		inf = open(infile, 'r')
		result = dict()
		for line in inf:
			if line == '\n':
				continue
			line = line[:-1].split(',')
			if valuetype == 'None':
				key = line[0]
				result[key] = 1
			else:
				key = line[0]
				value = line[1]
				if valuetype == 'int':
					result[key] = int(value)
				elif valuetype == 'str':
					result[key] = value
				else:
					None
		inf.close()
		return result


class JSHASH:
	"""
	Handle the hash of generated JS files.
	This class should be able to do the following:
	1. Get hash of JS from one link.
	2. Get hash of a list of JS and output the map to given file.
	3. Dump the generated hash to database.
	"""
	def path_checksum(self, paths):
		"""
		Recursively calculates a checksum representing the contents of all files
		found with a sequence of file and/or directory paths.
		"""
		if not hasattr(paths, '__iter__'):
			raise TypeError('sequence or iterable expected not %r!' % type(paths))

		def _update_checksum(checksum, dirname, filenames):
			for filename in sorted(filenames):
				path = path_join(dirname, filename)
				if isfile(path):
					fh = open(path, 'rb')
					while 1:
						buf = fh.read(4096)
						if not buf : break
						checksum.update(buf)
					fh.close()
		chksum = hashlib.sha1()

		for path in sorted([normpath(f) for f in paths]):
			if path_exists(path):
				if isdir(path):
					walk(path, _update_checksum, chksum)
				elif isfile(path):
					_update_checksum(chksum, dirname(path), basename(path))

			return chksum.hexdigest()


	def ls_path_checksum(self, directory, jslistfile = 'jslist.log', \
			 jshashfile = 'jshashs', jshash2countfile = 'jshash2count'):
		basedir = '../../data/TempData/LOGs/'
		jslistfile = basedir + jslistfile
		ls2file(directory, jslistfile)
		jshashfile = basedir + jshashfile
		jshash2countfile = basedir + jshash2countfile

		log = LOG(basedir + 'LOG')
		progress = log.STATUS(jslistfile)
		jslistf = open(jslistfile, 'r')
		count = 0

		jshashs = DictFile().file2dict(jshashfile, 'str')
		jshash2count = DictFile().file2dict(jshash2countfile, 'int')
		for jsdir in jslistf:
			count = count + 1
			if count <= progress:
				continue
			jsdir = jsdir[:-1]
			print 'Processing {0}'.format(jsdir)

			urlhash = jsdir.split('.')[-1]
			paths = [jsdir]
			jshash = self.path_checksum(paths)
			jshashs[urlhash] = jshash
			if jshash in jshash2count:
				jshash2count[jshash] = jshash2count[jshash] + 1
			else:
				jshash2count[jshash] = 1
			log.UPDATE(jslistfile, count)
		DictFile().dict2file(jshashs, jshashfile)
		DictFile().dict2file(jshash2count, jshash2countfile)

	def toDatabase(self, hashmap):
		con = psql.PsqlCon()
		for url_hash, js_hash in hashmap.items():
			con.psqlUpdateJSHash(url_hash, js_hash)



class LOG:
	"""
	Construct LOG information in LOG file.
	Each line is $filename,$progress
	progress here means the current processing line number of a given file.
	"""
	def __init__(self, logfile = '../../data/TempData/URLs/LOG'):
		self.logs = dict()
		self.logfile = logfile
		if not os.path.exists(self.logfile):
			# If the file doesn't exist, create one
			self.SAVELOG()
		self.logf = open(self.logfile, 'r')
		for line in self.logf:
			line = line[:-1].split(',')
			self.logs[line[0]] = int(line[1])
		self.logf.close()

	def LOADLOG(self):
		self.logf = open(self.logfile, 'r')
		self.logs = dict()
		for line in self.logf:
			line = line[:-1].split(',')
			self.logs[line[0]] = int(line[1])
		self.logf.close()

	def STATUS(self, filename):
		if filename not in self.logs:
			self.logs[filename] = 0
		return self.logs[filename]

	def UPDATE(self, filename, progress):
		self.logs[filename] = progress
		self.SAVELOG()
		self.LOADLOG()

	def SAVELOG(self):
		self.logf = open(self.logfile, 'w')
		for key in self.logs:
			self.logf.write(key + ',' + str(self.logs[key]) + '\n')
		self.logf.close()


class GetSample:
	def __init__(self, basedir = '/space/webcrawl-malware/data/TempData/LOGs/', \
			jslistfile = 'mjslist.log', jshashfile = 'mjshashs', jshash2countfile = 'mjshash2count'):
		self.jsAdd2TreeDir = '/space/webcrawl-malware/data/JSAdd2Tree/'
		self.basedir = basedir
		self.jslistfile = self.basedir + jslistfile
		self.jshashfile = self.basedir + jshashfile
		self.jshash2countfile = self.basedir + jshash2countfile
		self.jslist = DictFile().file2dict(self.jslistfile, 'None').keys()
		self.jslist = [item.split('/')[-1] for item in self.jslist]
		self.jshashs = DictFile().file2dict(self.jshashfile, 'str')
		self.jshash2count = DictFile().file2dict(self.jshash2countfile, 'int')

	def jslistNoTerry(self):
		self.jslist = [entry for entry in self.jslist if not entry.split('.')[0] == 'MTerry']

	def reInit(self, jslistfile = 'bjslist.log', jshashfile = 'bjshashs', jshash2countfile = 'bjshash2count'):
		self.jslistfile = self.basedir + jslistfile
		self.jshashfile = self.basedir + jshashfile
		self.jshash2countfile = self.basedir + jshash2countfile
		self.jslist = DictFile().file2dict(self.jslistfile, 'None').keys()
		self.jslist = [item.split('/')[-1] for item in self.jslist]
		self.jshashs = DictFile().file2dict(self.jshashfile, 'str')
		self.jshash2count = DictFile().file2dict(self.jshash2countfile, 'int')

	def isFailedUrl(self, dirname):
		dirname = '{0}/{1}/'.format(self.jsAdd2TreeDir, dirname)
		# the number of files in the directory is 23
		allCount = len(os.listdir(dirname))
		if not allCount == 23:
			return False
		# the number of files matching NativeEval* is 17
		nativeEvalCount = len([name for name in os.listdir(dirname) if fnmatch.fnmatch(name, 'NativeEval*')])
		if not nativeEvalCount == 17:
			return False
		# reloadUrl, failedUrl
		grepcmd = "grep 'failedUrl' {0}/3".format(dirname)
		if not commands.getoutput(grepcmd):
			return False 
		print dirname
		return True

	def isValid(self, dirname):
		urlhash = dirname.split('.')[1]
		if urlhash not in self.jshashs:
			return False
		return False if ( (self.jshash2count[self.jshashs[urlhash]] > 1) or self.isFailedUrl(dirname) ) else True

	def randomGet(self, count, label='M'):
		indexes = [i for i in range(0, len(self.jslist))]
		random.shuffle(indexes)
		progress = 0
		result = list()
		for index in indexes:
			if progress >= count:
				return result, True
			dirname = self.jslist[indexes[index]]
			if not dirname[0] == label:
				continue
			if self.isValid(dirname):
				result.append(dirname)
				progress = progress + 1
		if progress < count:
			return result, False
		else:
			return result, True

	def sortListByAlexa(self, alexaSortedListFile = '/space/webcrawl-malware/data/TempData/URLs/top-1m.csv'):
		ordereddict = DictFile().file2dict(alexaSortedListFile, 'str')
		self.sortedlist = [ ('B1.' + hashlib.sha256(ordereddict[ str(index) ]).hexdigest()) for index in range(1, len(ordereddict))]

	def sortListBySize(self):
		sortcmd = 'du -h {0}/B1.* | sort -h'.format(self.jsAdd2TreeDir)
		sizeSortedList = commands.getoutput(sortcmd).split('\n')
		self.sortedlist = [ line.split('\t')[-1].split('/')[-1] for line in sizeSortedList ]

	def topGet(self, count, benignsort):
		# If benignsort is Alexa, sort the list in alexa order,
		# the higher the rank, the higher the traffic flow;
		# If benignsort is Size, sort the list by size,
		# the higher the rank, the smaller the size.
		if benignsort == 'Alexa':
			self.sortListByAlexa()
		elif benignsort == 'Size':
			self.sortListBySize()

		progress = 0
		result = list()
		for dirname in self.sortedlist:
			if progress >= count:
				return result, True
			if self.isValid(dirname):
				result.append(dirname)
				progress = progress + 1
		if progress < count:
			return result, False
		else:
			return result, True

	def sampleList(self, maliciouscount = 500, benigncount = 500, outfile = 'samplelist', benignsort = 'Alexa'):
		self.outfile = self.basedir + outfile
		self.outf = open(self.outfile, 'w')

		maliciouslist, mstatus = self.randomGet(maliciouscount)

		self.reInit()
		print 'benignsort = {0}.............'.format(benignsort)
		if benignsort == 'Random':
			benignlist, bstatus = self.randomGet(benigncount, 'B')
		else:
			benignlist, bstatus = self.topGet(benigncount, benignsort)
		if bstatus and mstatus:
			print 'Successful, generating output.............'
			for item in benignlist:
				self.outf.write(item + '\n')
			for item in maliciouslist:
				self.outf.write(item + '\n')
		else:
			print 'No enough samples, malicious count {0}, benign count {1}'.format(len(maliciouslist), len(benignlist))

		self.outf.close()

	def copySampleList(self, samplefile = 'samplelist', jsdir = '/space/webcrawl-malware/data/JSAdd2Tree/', \
			outdir = '/space/webcrawl-malware/data/1000Samples/'):
		if not os.path.exists(outdir):
			os.makedirs(outdir)	

		self.inf = open(self.basedir + samplefile, 'r')
		for line in self.inf:
			if line == '\n':
				continue	
			dirname = line[:-1]
			print 'Processing {0}..........'.format(dirname)
			cpcmd = 'cp -r ' + jsdir + dirname + ' ' + outdir
			commands.getoutput(cpcmd)

	def copyTerry(self, jsdir = '/space/webcrawl-malware/data/JSAdd2Tree/',
			outdir = '/space/webcrawl-malware/data/1000SamplesWTerry/'):
		print 'Copying Terry examples.......'
		commands.getoutput('cp -r ' + jsdir + 'MTerry.* ' + outdir)


def main(bcount, mcount=None, wTerry=True, bSort='Alexa'):
	sample = GetSample()

	if mcount == None:
		samplefile = '{0}SamplesWTerryList'.format(bcount)
		outputdir = '/space/webcrawl-malware/data/{0}SamplesWTerry/'.format(bcount) 
		mcount = bcount / 2
		bcount = bcount - mcount
	else:
		samplefile = '{0}B{1}MSamplesWTerryList'.format(bcount, mcount)
		outputdir = '/space/webcrawl-malware/data/{0}B{1}MSamplesWTerry/'.format(bcount, mcount)

	# Check whether to include Terry's examples.
	if wTerry == True:
		sample.jslistNoTerry()
		mcount = mcount - 56
	else:
		samplefile = '{0}B{1}MSamplesList'.format(bcount, mcount)
		outputdir = '/space/webcrawl-malware/data/{0}B{1}MSamples/'.format(bcount, mcount)

	# Check whether bSort is valid
	if bSort != 'Alexa' and bSort != 'Size' and bSort != 'Random':
		print 'Unexpected Value for bSort!\nOnly Alexa and Size are Accepted.'
		sys.exit(1)

	# Generate sample list
	sample.sampleList(maliciouscount = mcount, benigncount = bcount, outfile = samplefile, benignsort = bSort)
	GetSample().copySampleList(samplefile = samplefile, outdir = outputdir)

	if wTerry == True:
		'''copySampleList will create directory if not exist'''
		GetSample().copyTerry(jsdir ='/space/webcrawl-malware/data/JSAdd2Tree/', outdir = outputdir)


if __name__ == '__main__':
	# print path_checksum([r'/tmp', '/etc/hosts'])
	# thepath = '/home/ruian/chromium/webcrawl-malware/data/JSAdd2Tree/MTerry.becd2b6920a629ed9eece578c6eef2ce2140fd60f7976fb308d9f13c9114678a'
	# print path_checksum(['/home/ruian/chromium/webcrawl-malware/data/JSAdd2Tree/MTerry.becd2b6920a629ed9eece578c6eef2ce2140fd60f7976fb308d9f13c9114678a'])

	'''
	Get the jshash for benign sites and malicious sites separately.
	'''
	# JSHASH().ls_path_checksum('../../data/JSAdd2Tree/M*', 'mjslist.log', 'mjshashs', 'mjshash2count')
	# JSHASH().ls_path_checksum('../../data/JSAdd2Tree/B*', 'bjslist.log', 'bjshashs', 'bjshash2count')

	# Select stuff needed to process
	'''
	# 1000Samples
	GetSample().sampleList(oufile = '1000SamplesList')
	GetSample().copySampleList()
	'''

	# Generate Samples including Terry's examples
	# main(20000)
	if len(sys.argv) == 3:
		B = int(sys.argv[1])
		M = int(sys.argv[2])
		main(B, M)
	elif len(sys.argv) == 4:
		B = int(sys.argv[1])
		M = int(sys.argv[2])
		bSort = sys.argv[3]
		main(bcount = B, mcount = M, bSort = bSort)
	else:
		main(300, 3000)
	# main(1000, 10000, False)

